package xyz.syyrjx.blog.util.ai;

import xyz.syyrjx.blog._enum.State;
import xyz.syyrjx.blog.dao.ai.AiIdentityDao;
import xyz.syyrjx.blog.dao.ai.AiIdentityGradeDao;
import xyz.syyrjx.blog.entity.ai.AiIdentity;
import xyz.syyrjx.blog.entity.ai.AiIdentityGrade;
import xyz.syyrjx.blog.exception.AiException;
import xyz.syyrjx.blog.util.CachePool;
import xyz.syyrjx.blog.util.ConstValuePool;

import javax.servlet.http.HttpSession;
import java.util.Map;

/**
 * @author 大地崩坏苍蝇兽
 * @date 2024/2/20 18:16
 * @description ...
 * @since v0.0.1
 */
public class AiUseCheckUtil {

    public static AiIdentityDao aiIdentityDao;
    public static AiIdentityGradeDao aiIdentityGradeDao;

    /**
     * 返回两个cookie，第一个记录用量，第二个记录id
     * @param session
     * @param countModel
     * @return
     */
    public static void usageCheck(String aiIdentity, String countModel
            , String usageCountFlag, int maxUsageCount) {
        // 查询ai对话次数记录表
        Map<String, Integer> countMap = CachePool.AI_USAGE_MAP.get(aiIdentity);
        if (countMap == null) {
            throw new AiException("没有查询到您的Ai对话记录表，请尝试刷新页面获取记录表");
        }
        // 选中模型对话次数
        Integer usageCount = countMap.getOrDefault(countModel, 0);

        // 放行等级
        AiIdentity aiIdentityObj = aiIdentityDao
                .selectAiIdentity(
                        new AiIdentity(null, aiIdentity, null, State.NORMAL_STATE.getCode())
                );
        if (aiIdentityObj == null) {
            throw new AiException("身份信息有误");
        }
        AiIdentityGrade aiIdentityGrade = aiIdentityGradeDao
                .selectAiIdentityGrade(
                        new AiIdentityGrade(null, aiIdentityObj.getId(), null, State.NORMAL_STATE.getCode())
                );
        // 如果被放行，修改最大值
        if (aiIdentityGrade != null) {
            maxUsageCount = CachePool.aiIdGradeCountMap
                    // 取模型记录
                    .get(usageCountFlag)
                    // 取等级记录
                    .get(aiIdentityGrade.getGrade());
        }
        // 超额
        if (usageCount >= maxUsageCount) {
            throw new AiException("您对本模型只能进行" + maxUsageCount + "轮对话，如有需要请联系管理员") ;
        }
        countMap.put(countModel, usageCount + 1);
    }
}
